Skip to content

[Core] Fix L40 cublas heuristics - #3440

Open
jberchtold-nvidia wants to merge 5 commits into
NVIDIA:mainfrom
jberchtold-nvidia:jberchtold/fix-l40-cublas-heuristic
Open

[Core] Fix L40 cublas heuristics#3440
jberchtold-nvidia wants to merge 5 commits into
NVIDIA:mainfrom
jberchtold-nvidia:jberchtold/fix-l40-cublas-heuristic

Conversation

@jberchtold-nvidia

Copy link
Copy Markdown
Collaborator

Description

In certain cases on L40 we would allow a cuBLASLt heuristic that performed split-K with partial results stored in BF16. This lead to some loss of precision in the results. This PR fixes this by disallowing such heuristics

H100 and above were not observed to be affected by this.

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring

Changes

  • Disallow split-K BF16 cuBLASLt heuristics

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

@jberchtold-nvidia

Copy link
Copy Markdown
Collaborator Author

/te-ci L1

@greptile-apps

greptile-apps Bot commented Aug 28, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

The PR constrains cuBLASLt heuristic selection to prevent reduced-precision split-K reductions.

  • Sets the reduction-scheme preference mask to retain compute-type reductions.
  • Prevents split-K partial results from being converted to BF16 before final accumulation.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename Overview
transformer_engine/common/gemm/cublaslt_gemm.cu Adds a cuBLASLt preference that excludes output-type split-K reductions while retaining non-split and compute-type reduction algorithms.

Reviews (2): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile

Comment on lines +755 to +756
if (A_type == CUDA_R_16BF && B_type == CUDA_R_16BF && D_type == CUDA_R_16BF &&
gemm_compute_type == CUBLAS_COMPUTE_32F) {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why do we only do this for the BF16 case? I think we should just set this always.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good point, I've generalized it to always set this

@ptrendx ptrendx added the 2.20 label Aug 31, 2026
@ptrendx

ptrendx commented Aug 31, 2026

Copy link
Copy Markdown
Member

/te-ci

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants